# 音乐下载器模块

import os
import sys
# 将项目根目录添加到Python搜索路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import requests
import time
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor, as_completed
from config import settings
from utils import setup_logger, create_dir_if_not_exists, clean_filename, retry_on_error, write_log_to_file, format_file_size

logger = setup_logger('downloader')


class MusicDownloader:
    """音乐下载器类"""
    def __init__(self):
        self.save_path = settings.MUSIC_SAVE_PATH
        self.max_retries = settings.MAX_RETRY
        self.timeout = settings.TIMEOUT
        self.max_concurrent = settings.MAX_CONCURRENT
        self.headers = settings.HEADERS.copy()
        
        # 确保保存目录存在
        create_dir_if_not_exists(self.save_path)

    @retry_on_error()
    def download_file(self, url, file_path, show_progress=True):
        """下载单个文件"""
        try:
            # 检查文件是否已存在
            if os.path.exists(file_path):
                logger.info(f"文件已存在，跳过下载: {file_path}")
                return True, file_path, os.path.getsize(file_path)
            
            # 根据URL域名设置不同的请求头
            headers = self.headers.copy()
            if 'kuwo.cn' in url:
                # 使用用户提供的Firefox浏览器请求头配置kuwo.cn
                headers = {
                    "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:137.0) Gecko/20100101 Firefox/137.0",
                    "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
                    "Accept-Language": "zh-CN,zh;q=0.8,zh-TW;q=0.7,zh-HK;q=0.5,en-US;q=0.3,en;q=0.2",
                    "Accept-Encoding": "gzip, deflate, br, zstd",
                    "Connection": "keep-alive",
                    "Cookie": "Hm_lvt_cdb524f42f0ce19b169a8071123a4797=1759917433; Hm_lpvt_cdb524f42f0ce19b169a8071123a4797=1759917433; HMACCOUNT=A6B64871CBEA241D; _ga=GA1.2.90125026.1759917433; _gid=GA1.2.163174485.1759917433; _ga_ETPBRPM9ML=GS2.2.s1759917433$o1$g1$t1759917439$j54$l0$h0",
                    "Upgrade-Insecure-Requests": "1",
                    "Sec-Fetch-Dest": "document",
                    "Sec-Fetch-Mode": "navigate",
                    "Sec-Fetch-Site": "none",
                    "Sec-Fetch-User": "?1",
                    "Priority": "u=0, i",
                    "Pragma": "no-cache",
                    "Cache-Control": "no-cache"
                }
            
            # 发送请求
            response = requests.get(url, headers=headers, timeout=self.timeout, stream=True)
            response.raise_for_status()
            
            # 获取文件大小
            file_size = int(response.headers.get('content-length', 0))
            
            # 创建目录（如果需要）
            file_dir = os.path.dirname(file_path)
            create_dir_if_not_exists(file_dir)
            
            # 下载文件
            with open(file_path, 'wb') as f:
                if show_progress and file_size > 0:
                    # 使用tqdm显示下载进度
                    with tqdm(total=file_size, unit='B', unit_scale=True, desc=os.path.basename(file_path)) as pbar:
                        for chunk in response.iter_content(chunk_size=8192):
                            if chunk:
                                f.write(chunk)
                                pbar.update(len(chunk))
                else:
                    # 不显示进度
                    for chunk in response.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
            
            # 检查下载文件大小
            downloaded_size = os.path.getsize(file_path)
            if file_size > 0 and downloaded_size < file_size * 0.95:
                # 文件可能不完整
                logger.warning(f"文件可能不完整: {file_path}, 预期: {format_file_size(file_size)}, 实际: {format_file_size(downloaded_size)}")
                os.remove(file_path)  # 删除不完整的文件
                return False, file_path, 0
            
            logger.info(f"文件下载成功: {file_path}, 大小: {format_file_size(downloaded_size)}")
            write_log_to_file(f"下载成功: {url} -> {file_path}", "INFO")
            return True, file_path, downloaded_size
        except requests.exceptions.RequestException as e:
            logger.error(f"下载请求失败: {url}, 错误: {e}")
            write_log_to_file(f"下载失败 (请求错误): {url}, 错误: {e}", "ERROR")
            # 清理不完整的文件
            if os.path.exists(file_path):
                try:
                    os.remove(file_path)
                except:
                    pass
            return False, file_path, 0
        except Exception as e:
            logger.error(f"下载文件失败: {url}, 错误: {e}")
            write_log_to_file(f"下载失败 (未知错误): {url}, 错误: {e}", "ERROR")
            # 清理不完整的文件
            if os.path.exists(file_path):
                try:
                    os.remove(file_path)
                except:
                    pass
            return False, file_path, 0

    def download_music(self, song_data, artist_dir=True, album_dir=False):
        """下载音乐文件，支持按歌手和专辑分类保存"""
        try:
            # 提取歌曲信息
            song_id = song_data.get('song_id', '')
            title = song_data.get('title', '未知歌曲')
            artist = song_data.get('artist', '未知歌手')
            album = song_data.get('album', '')
            url = song_data.get('url', '')
            
            if not url:
                logger.error(f"没有找到下载链接: {title} - {artist}")
                return False, None, 0
            
            # 构建保存路径
            save_dir = self.save_path
            
            # 按歌手分类
            if artist_dir:
                artist = clean_filename(artist)
                save_dir = os.path.join(save_dir, artist)
            
            # 按专辑分类
            if album_dir and album:
                album = clean_filename(album)
                save_dir = os.path.join(save_dir, album)
            
            # 从URL中提取文件扩展名
            # 先处理URL参数，获取主URL部分
            main_url = url.split('?')[0].lower()
            
            # 使用字符串处理方法提取扩展名
            last_dot_index = main_url.rfind('.')
            if last_dot_index != -1 and last_dot_index < len(main_url) - 1:
                url_ext = main_url[last_dot_index:]
                # 确保扩展名合理，不超过5个字符
                if len(url_ext) > 5:
                    url_ext = '.mp3'
            else:
                # 没有找到有效扩展名，使用默认值
                url_ext = '.mp3'
            
            # 构建文件名，使用URL中的扩展名
            filename = clean_filename(f"{title} - {artist}{url_ext}")
            file_path = os.path.join(save_dir, filename)
            
            # 下载文件
            success, file_path, file_size = self.download_file(url, file_path)
            
            return success, file_path, file_size
        except Exception as e:
            logger.error(f"下载音乐失败: {song_data.get('title', '未知歌曲')}, 错误: {e}")
            return False, None, 0

    def batch_download(self, songs_data, artist_dir=True, album_dir=False, max_workers=None):
        """批量下载音乐"""
        if not songs_data:
            logger.info("没有要下载的歌曲")
            return [], []
        
        # 设置最大工作线程数
        if max_workers is None:
            max_workers = min(self.max_concurrent, len(songs_data))
        
        logger.info(f"开始批量下载，共 {len(songs_data)} 首歌曲，最大并发数: {max_workers}")
        write_log_to_file(f"开始批量下载，共 {len(songs_data)} 首歌曲", "INFO")
        
        success_list = []  # 成功下载的歌曲
        failed_list = []   # 下载失败的歌曲
        
        # 使用线程池进行并发下载
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            # 提交所有下载任务
            future_to_song = {
                executor.submit(self.download_music, song, artist_dir, album_dir): song 
                for song in songs_data
            }
            
            # 处理完成的任务
            for future in as_completed(future_to_song):
                song = future_to_song[future]
                try:
                    success, file_path, file_size = future.result()
                    if success:
                        # 更新歌曲信息
                        song['file_path'] = file_path
                        song['file_size'] = file_size
                        success_list.append(song)
                        logger.info(f"成功下载: {song.get('title')} - {song.get('artist')}")
                    else:
                        failed_list.append(song)
                        logger.warning(f"下载失败: {song.get('title')} - {song.get('artist')}")
                except Exception as e:
                    failed_list.append(song)
                    logger.error(f"处理下载任务异常: {song.get('title')} - {song.get('artist')}, 错误: {e}")
        
        total = len(songs_data)
        success_count = len(success_list)
        failed_count = len(failed_list)
        
        logger.info(f"批量下载完成，总数: {total}, 成功: {success_count}, 失败: {failed_count}")
        write_log_to_file(f"批量下载完成，总数: {total}, 成功: {success_count}, 失败: {failed_count}", "INFO")
        
        return success_list, failed_list

    def check_url_validity(self, url):
        """检查URL是否有效"""
        try:
            response = requests.head(url, headers=self.headers, timeout=self.timeout, allow_redirects=True)
            # 检查状态码和内容类型
            if response.status_code == 200 and 'audio' in response.headers.get('content-type', ''):
                return True, response.headers.get('content-length', 0)
            return False, 0
        except Exception as e:
            logger.error(f"检查URL有效性失败: {url}, 错误: {e}")
            return False, 0

    def get_download_speed(self, url, test_duration=5):
        """测试下载速度"""
        try:
            start_time = time.time()
            start_size = 0
            
            response = requests.get(url, headers=self.headers, timeout=self.timeout, stream=True)
            response.raise_for_status()
            
            # 测试下载一小段数据
            with open(os.devnull, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
                        start_size += len(chunk)
                        
                    # 超过测试时间则停止
                    if time.time() - start_time > test_duration:
                        break
            
            end_time = time.time()
            duration = end_time - start_time
            
            if duration > 0:
                speed = start_size / duration  # 字节/秒
                return True, speed
            return False, 0
        except Exception as e:
            logger.error(f"测试下载速度失败: {url}, 错误: {e}")
            return False, 0

    def set_proxy(self, proxy=None):
        """设置代理服务器"""
        self.proxy = proxy

    def set_headers(self, headers):
        """设置请求头"""
        if headers and isinstance(headers, dict):
            self.headers.update(headers)
            logger.info("请求头已更新")